import time
import black
import operator
import functools
from typing import Annotated, TypedDict, Sequence, Literal

from src.module.TongyiModel import TongyiModel

from langchain_experimental.utilities import PythonREPL

from langchain_core.messages import ToolMessage, BaseMessage, AIMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langgraph.graph import StateGraph, START, END
from langchain.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langgraph.prebuilt import ToolNode

run_code = PythonREPL()


class Tools:
    @staticmethod
    @tool
    def exec_python_code(code: Annotated[str, "需要执行的 python 代码"]):
        """ 使用这个工具来执行Python代码。代码应该是完整的，可执行的全部代码。而不是某一部分的代码片段。
            如果你想查看某个值的输出，应该使用print(...)，这个输出对用户可见。
        """
        try:
            result = run_code.run(code)
            time.sleep(10)
        except BaseException as e:
            return f"执行失败。错误: {repr(e)}"

        return  f"成功执行:\n```python\n{code}\n```\nStdout: {result} \n\n如果你已完成所有任务，请回复FINAL ANSWER。"

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    sender: str # 追踪当前的节点是谁得出的答案

class LangGraphMoreAgent(TongyiModel):
    ## 将函数作为节点的辅助函数
    ## 为什么要将函数通过 functools.partial 创建为 节点，agent 本身不就可以当作节点使用吗？
    ## 因为这里的场景使用到了自定义节点的内容输出以及节点的 sender 标识，因此，agent 原始输出不满足这一点，所以才会将 agent 在包装一层。
    @staticmethod
    def create_func_to_node(state, agent, name):
        result = agent.invoke(state)

        if not isinstance(result, ToolMessage):
            # 将输出的内容进行重新包装，主要是为了 name 参数，也就是上面定义的 StateMessage 中的 sender
            result = AIMessage(**result.dict(exclude={"type", "name"}), name=name)

        return { "messages": [result], "sender": name }

    ## 创建 agent
    def create_agent(self, system_message: str, tools: list):
        prompt = ChatPromptTemplate.from_messages([
            ('system', """
                你是⼀个有帮手的AI助⼿，与其他助⼿合作。
                使⽤提供的⼯具来推进问题的回答。
                如果你不能完全回答，没关系，另⼀个拥有不同⼯具的助⼿会接着你的位置继续帮助。
                执⾏你能做的以取得进展。
                如果你或其他助⼿有最终答案或交付物，
                在你的回答前加上FINAL ANSWER，以便团队知道停⽌。
                你可以使⽤以下⼯具: {tool_names}。\n{system_message}
            """),
            MessagesPlaceholder(variable_name='messages')
        ])

        # partial 替换内容，将消息内容中的占位符替换为 指定的内容
        prompt = prompt.partial(system_message=system_message)
        prompt = prompt.partial(tool_names=",".join([item.name for item in tools]))
        return prompt | self.model.bind_tools(tools)

    ## 路由函数，自定义 graph 的走向
    @staticmethod
    def router(state: AgentState) -> Literal['__end__', 'call_tool', 'continue']:
        messages = state['messages'][-1]
        if not messages:
            return 'continue'

        if isinstance(messages.tool_calls, list) and len(messages.tool_calls) != 0:
            # 调用工具
            return 'call_tool'
        elif 'FINAL ANSWER' in messages.content:
            # 在 agent 函数中，prompt 中提到了，当返回的内容中包含：FINAL ANSWER 就表示程序执行完成了
            return '__end__'
        else:
            # 继续执行 node 节点
            return 'continue'

    def start(self):
        ## 将工具作为节点装载为节点
        search_tool, chat_generator_tool = TavilySearchResults(max_result=5), Tools().exec_python_code
        tools = [search_tool, chat_generator_tool]
        tools_node = ToolNode(tools)

        ## 创建绘图的 agent
        chart_generator_agent = self.create_agent(system_message='', tools=[chat_generator_tool])
        chart_node = functools.partial(self.create_func_to_node, agent=chart_generator_agent, name=chat_generator_tool.name)

        ## 创建查询数据的 agent
        search_agent = self.create_agent(system_message=f'你应该提供精准的数据给到{chat_generator_tool.name}使用', tools=[search_tool])
        search_node = functools.partial(self.create_func_to_node, agent=search_agent, name='search_node')

        ## 添加节点
        graph_builder = StateGraph(AgentState)
        graph_builder.add_node('search_node', search_node) ## 查询数据
        graph_builder.add_node('chart_node', chart_node) ## 生成图表
        graph_builder.add_node('tools_node', tools_node) ## 工具节点

        ## 创建 graph，全部用条件边路由来决定流程的走向。
        graph_builder.add_edge(START, "search_node")
        graph_builder.add_conditional_edges(
            'search_node',
            self.router,
            { 'continue': 'chart_node', 'call_tool': 'tools_node', '__end__': END }
        )
        graph_builder.add_conditional_edges(
            'chart_node',
            self.router,
            {'continue': 'search_node', 'call_tool': 'tools_node', '__end__': END} # 当发现仍然需要 continue 那么就重新执行 search，重置流程
        )
        graph_builder.add_conditional_edges(
            'tools_node',
            lambda x: x['sender'], ## 匿名函数，可以理解为 箭头函数，x 表示的就是 AgentState 参数
            {
                "search_node": "search_node",
                "chart_node": "chart_node",
            }
        )
        graph = graph_builder.compile()

        #事件流
        events = graph.stream(
            {
                "messages": [
                    HumanMessage(
                        content="获取过去5年AI软件市场规模，"
                                " 然后绘制一条折线图。"
                                " 一旦你编写好代码，完成任务。"
                    )
                ],
            },
            # 图中最多执行的步骤数
            {"recursion_limit": 150},
        )
        # 打印事件流中的每个状态
        for s in events:
            print(f"\n\n {s} \n\n----")

